# -*- coding: utf-8 -*-
"""
问答链实现模块

本模块实现了基于LangChain框架的问答处理链，整合了检索器和大语言模型，
提供端到端的问答解决方案。

核心功能：
1. 问题理解：接收用户输入的自然语言问题
2. 信息检索：使用混合检索器查找相关上下文
3. 答案生成：基于检索到的上下文生成准确答案
4. 结果处理：格式化和优化输出结果

设计特点：
- 模块化设计：检索器和LLM可独立配置和替换
- 可扩展性：支持自定义提示模板和输出解析器
- 灵活性：可适应不同领域的问答需求
- 高效性：通过检索减少无关信息，提高生成质量

适用场景：
- 知识库问答系统
- 文档检索问答
- 客服机器人
- 教育辅助系统
"""
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from typing import Optional, Dict, Any
from loguru import logger

from src.core.llm.ollama_client import OllamaClient
from src.core.retrieval.hybrid_retriever import HybridRetriever

class QAChain:
    """
    问答处理链
    
    这是一个完整的问答处理流水线，整合了信息检索和文本生成能力。
    采用RAG（Retrieval-Augmented Generation）架构，先检索相关信息，
    再基于检索结果生成答案。
    
    核心组件：
    1. LLM客户端：负责文本生成和理解
    2. 检索器：负责从知识库中查找相关信息
    3. 提示模板：定义LLM的输入格式
    4. 输出解析器：处理和格式化LLM的输出
    
    工作流程：
    问题输入 -> 信息检索 -> 上下文构建 -> LLM生成 -> 结果解析 -> 答案输出
    
    Attributes:
        llm_client (OllamaClient): 大语言模型客户端实例
        retriever (HybridRetriever): 混合检索器实例
        output_parser (StrOutputParser): 输出解析器
        prompt_template (PromptTemplate): 提示模板
    """
    
    def __init__(self, llm_client: OllamaClient, retriever: HybridRetriever):
        """
        初始化问答链
        
        Args:
            llm_client (OllamaClient): 已配置的Ollama客户端实例
            retriever (HybridRetriever): 已配置的混合检索器实例
        
        Raises:
            ValueError: 当传入的参数为None时抛出
        
        Example:
            >>> from src.core.llm.ollama_client import OllamaClient
            >>> from src.core.retrieval.hybrid_retriever import HybridRetriever
            >>> 
            >>> llm = OllamaClient("deepseek-r1:32b")
            >>> retriever = HybridRetriever(qa_data)
            >>> qa_chain = QAChain(llm, retriever)
        """
        if llm_client is None:
            raise ValueError("LLM客户端不能为空")
        if retriever is None:
            raise ValueError("检索器不能为空")
            
        self.llm_client = llm_client
        self.retriever = retriever
        self.output_parser = StrOutputParser()
        
        # 默认提示模板
        # 采用结构化的提示格式，明确指定上下文和问题的位置
        # 这种格式有助于模型更好地理解任务要求
        self.prompt_template = PromptTemplate(
            input_variables=["question", "context"],
            template="""
            基于以下上下文回答问题：
            
            上下文：{context}
            
            问题：{question}
            
            请提供准确、简洁的答案：
            """
        )
        
        logger.info(f"问答链初始化完成，使用模型: {llm_client.model_name}")
        self.conversation_history = []  # 添加会话历史
    
    def run(self, question: str, use_history: bool = True) -> str:
        """
        运行问答链，处理用户问题并生成答案
        
        这是问答链的核心方法，执行完整的问答流程：
        1. 输入验证：检查问题格式和内容
        2. 信息检索：使用检索器查找相关上下文
        3. 提示构建：将问题和上下文组合成LLM输入
        4. 答案生成：调用LLM生成回答
        5. 结果处理：解析和格式化最终答案
        
        Args:
            question (str): 用户输入的问题文本
        
        Returns:
            str: 生成的答案文本
        
        Raises:
            ValueError: 当问题为空或格式不正确时
            RuntimeError: 当检索或生成过程出现错误时
        
        Example:
            >>> qa_chain = QAChain(llm_client, retriever)
            >>> answer = qa_chain.run("什么是人工智能？")
            >>> print(answer)
            人工智能是计算机科学的一个分支...
        """
        # 输入验证
        if not question or not question.strip():
            raise ValueError("问题不能为空")
        
        question = question.strip()
        logger.info(f"开始处理问题: {question[:50]}...")
        
        try:
            # 检索相关信息
            # 使用混合检索器查找最相关的上下文信息
            context = self.retriever.get_best_answer(question)
            logger.debug(f"检索到上下文: {str(context)[:100]}...")
            
            # 构建提示
            # 将问题和检索到的上下文组合成结构化的提示
            prompt = self.prompt_template.format(
                question=question,
                context=context
            )
            logger.debug(f"构建的提示长度: {len(prompt)} 字符")
            
            # 构建包含历史的上下文
            if use_history and self.conversation_history:
                history_context = "\n".join([
                    f"Q: {item['question']}\nA: {item['answer']}"
                    for item in self.conversation_history[-3:]  # 保留最近3轮对话
                ])
                context = f"对话历史：\n{history_context}\n\n当前上下文：{context}"
            
            # 生成答案
            # 调用LLM生成基于上下文的答案
            answer = self.llm_client.generate(prompt)
            logger.debug(f"生成的原始答案长度: {len(answer)} 字符")
            
            # 解析和格式化输出
            final_answer = self.output_parser.parse(answer)
            logger.info(f"问答处理完成，答案长度: {len(final_answer)} 字符")
            
            # 保存到历史
            if use_history:
                self.conversation_history.append({
                    "question": question,
                    "answer": final_answer
                })
            
            return final_answer
            
        except Exception as e:
            logger.error(f"问答处理过程中出现错误: {str(e)}")
            raise RuntimeError(f"问答处理失败: {str(e)}")
    
    def set_prompt_template(self, template: str, input_variables: list = None) -> None:
        """
        设置自定义提示模板
        
        允许用户根据特定需求自定义提示模板，以适应不同的问答场景。
        
        Args:
            template (str): 新的提示模板字符串
            input_variables (list, optional): 模板中的变量列表，默认为["question", "context"]
        
        Example:
            >>> template = "请根据以下信息回答问题：\n信息：{context}\n问题：{question}\n答案："
            >>> qa_chain.set_prompt_template(template)
        """
        if input_variables is None:
            input_variables = ["question", "context"]
            
        self.prompt_template = PromptTemplate(
            input_variables=input_variables,
            template=template
        )
        logger.info("提示模板已更新")
    
    def get_chain_info(self) -> Dict[str, Any]:
        """
        获取问答链的配置信息
        
        Returns:
            Dict[str, Any]: 包含链配置信息的字典
        """
        return {
            "llm_model": self.llm_client.model_name,
            "retriever_type": type(self.retriever).__name__,
            "prompt_template": self.prompt_template.template,
            "input_variables": self.prompt_template.input_variables
        }